from multiprocessing import Pool
import time
from mechine import Heat_recovery_system, steam_turbine, Condenser, D_pre, zero_pre
import numpy as np
import matplotlib.pyplot as plt
import os
import math
from param_ident2 import share_data
import pickle

plt.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体：解决plot不能显示中文问题
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题
# 计算净功率
def Net_power(Ne, n_opt, n_fan=56):
    power_consumption_fan = (n_opt  /110) ** 3 * 115 * 0.001 * n_fan
    current_net_power = Ne - power_consumption_fan
    return power_consumption_fan, current_net_power

def sim_run(simtime, data, h=0.01, sample_time=5, state_result=None, ita_current=0.9, kc_current=0.04, list=[], nelist=[]):
    if state_result is None:
        state_result = []

    turbine = steam_turbine(ita_current)
    heatsystem = Heat_recovery_system()
    coldsystem = Condenser()

    Ne, Ps1, Ps2, P1, P2, T1, T2, Tj, Ty2, Ps3, \
        Ps4, Ps5, Ps6, Ps7, Tn7, Tl7, Tw7, hd7, Dd6, hd6, Tn6, \
        Tl6, Tw6, Dd5, hd5, Tn5, Tl5, Tw5, Dd3, hd3, Tn3, Tl3, Tw3, \
        Dd2, hd2, Tn2, Tl2, Tw2, Dd1, hd1, Tn1, Tl1, Tw1, Twc_m, Tw4, \
        drhol4_dt, rhol4_m, Pb = (state_result[item] for item in range(len(state_result)))

    # 自找数据
    Ne_DT = data[1]  # 机组负荷
    Pt_DT = data[2]  # 主汽压力
    Dsup_DT = data[3]  # 主汽流量
    Pre_real = data[4]  # 再热蒸汽压力
    Re_wat_T = data[5]  # 再热蒸汽温度
    Cond_wat_T = data[6]  # 凝结水温度
    Cond_wat_D = data[7]  # 凝结水流量
    Feed_wat_D = data[8]  # 主给水流量
    Sup_wat_T = data[9]  # 主汽温度
    Ta = data[10]  # 环境温度
    # n_num = data[11]  # 风机运行台数
    n_opt_initial = data[11]  # 风机转速, 二分法数据没有风机台数
    # Pb_DT = data[-1] / 1000  # 背压
    Re_attw_D = 0
    # if Pb != 0.015:
    #     Pb_DT = Pb
    #Pb_DT = data[-1]
    Pb_DT = Pb

    for j in range(int(sample_time / h)):
        # 汽轮机仿真
        Ds, pt, ht, D1, p1, D1j, h1, D2, p2, Dj2, h2, P2, T2, D2r, D3, p3, h3, Dj3, D4, p4, h4, \
            Dj4, D5, p5, h5, Dj5, D6, p6, h6, Dj6, D7, p7, h7, Dj7, hc, Dc, P1, T1, Tj, Ty2, Ne, Cond_T, Ne8 = turbine.simulation(
            Sup_wat_T, Pt_DT, Dsup_DT, Ps1, Ps2, Re_attw_D,
            Ne_DT, P1, P2, Pre_real, T1, T2, Re_wat_T, Tj, Ty2,
            Ps3, Ps4, Ps5, Ps6, Ps7, Pb, h)

        # 凝汽器仿真
        _, Pb, Dec, Tc, Hc1, Vf, NTU, Qc = coldsystem.simulation(hc, Dc, Pb, n_opt_initial, Ta, h, kc_current)

        # 回热系统仿真
        Ps7, Dd7, hd7, Tn7, Tl7, Tw7, Ps6, Dd6, hd6, Tn6, Tl6, Tw6, Ps5, Dd5, hd5, Tn5, Tl5, Tw5,\
        Ps4, Tw4, Ps3, Dd3, hd3, Tn3, Tl3, Tw3, Ps2, Dd2, hd2, Tn2, Tl2, Tw2, Ps1, Dd1, hd1, Tn1, Tl1, Tw1,\
        rhol4, Twc_m, Tw4_m, rhol4_m, drhol4_dt = heatsystem.simulation(
            Cond_wat_T, Tn7, Tl7, Ps7, Tw7, D7, h7, hd7, Dd6, hd6,Tn6, Tl6, Ps6, Tw6, D6, h6, Dd5, hd5,Tn5, Tl5, Ps5,
            Tw5, D5, h5, Ps4, Dd3, hd3, D4, h4, Dsup_DT, Tn3, Tl3, Ps3, Tw3, D3, h3, Dd2, hd2,Tn2, Tl2, Ps2, Tw2,
            D2, h2, Dd1, hd1, Tn1, Tl1, Ps1, Tw1, D1, h1,Cond_wat_D, Feed_wat_D, Tw4,Twc_m,drhol4_dt,rhol4_m, h)
        list.append(Pb)
        nelist.append(Ne)

    state_result = [Ne, Ps1, Ps2, P1, P2, T1, T2, Tj, Ty2, Ps3,
                    Ps4, Ps5, Ps6, Ps7, Tn7, Tl7, Tw7, hd7, Dd6, hd6,
                    Tn6, Tl6, Tw6, Dd5, hd5, Tn5, Tl5, Tw5, Dd3, hd3,
                    Tn3, Tl3, Tw3, Dd2, hd2, Tn2, Tl2, Tw2, Dd1, hd1,
                    Tn1, Tl1, Tw1, Twc_m, Tw4, drhol4_dt, rhol4_m, Pb]

    mid_vars = [Vf, NTU, Qc, Dec]


    return state_result, mid_vars, list, nelist

def optimize_for_speed(data1, m1, simtime, h, sample_time, tolerance, ita_current, kc_current, stable_state=None):
    #data1[11] = m1
    print(f'优化输入ita:{ita_current}和kc:{kc_current}')
    if stable_state is None:
        state_result = [330, 3.571, 2.261, 2.1, 2, 480, 500, 520, 900, 1.203,
                        0.6006, 0.3668, 0.1441, 0.04862, 62, 70, 77.85, 253, 5, 350,
                        83, 95, 107, 5, 474, 113, 120, 138, 5, 712,
                        167, 180, 188, 5, 823, 200, 210, 218, 5, 963,
                        220, 236, 245, 54, 180, 0.01, 0.51, 0.00975]
        # pb_list = [0.00975]
        # ne_list = [330]
        print(f'初始状态:{state_result}')
    else:
        state_result = stable_state
        # pb_list = [state_result[-1]]
        # #data1[1] = stable_state[0]
        # #data1[-1] = state_result[-1]
        # # print(f'输入数据：{data1}')
        # ne_list = [stable_state[0]]
    pb_list = []
    ne_list = []

    turbine = steam_turbine(ita_current)
    heatsystem = Heat_recovery_system()
    coldsystem = Condenser()
    Dec = 0
    for one_data in data1:
        one_data = one_data[1:]
        non_zeros = [0, 1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13, 14, 17, 18, 19, 22, 23]
        if np.any(one_data[non_zeros] == 0):
            continue
        # 解包当前状态
        (Ne, Ps1, Ps2, P1, P2, T1, T2, Tj, Ty2, Ps3, Ps4, Ps5, Ps6, Ps7,
         Tn7, Tl7, Tw7, hd7, Dd6, hd6, Tn6, Tl6, Tw6, Dd5, hd5, Tn5,
         Tl5, Tw5, Dd3, hd3, Tn3, Tl3, Tw3, Dd2, hd2, Tn2, Tl2, Tw2,
         Dd1, hd1, Tn1, Tl1, Tw1, Twc_m, Tw4, drhol4_dt, rhol4_m, Pc_sim
         ) = state_result

        # 读取当前时间步真实数据
        Ne_DT = one_data[0]  # 真实汽轮机负荷
        Pc_DT = one_data[17]  # 背压（单位：KPa）
        Pt_DT = one_data[3]
        Dsup_DT = one_data[1]
        Pre_real = one_data[5]
        Re_wat_T = one_data[6]
        Cond_wat_T = one_data[12]
        Cond_wat_D = one_data[11]
        Feed_wat_D = one_data[10]
        Sup_wat_T = one_data[2]
        n = m1
        Ta = one_data[18]
        # print(f'Ne_DT:{Ne_DT},Pc_DT{Pc_DT},Pt_DT{Pt_DT},Dsup_DT{Dsup_DT},Pre_real{Pre_real},Re_wat_T{Re_wat_T},Cond_wat_T{Cond_wat_T},Cond_wat_D{Cond_wat_D},Feed_wat_D{Feed_wat_D},Sup_wat_T{Sup_wat_T},Ta{Ta}')

        ne_list.append(Ne)
        pb_list.append(Pc_sim * 1000)
        state_result_new = []
        # 内层仿真：调用汽轮机和凝汽器模拟
        for j in range(int(sample_time / h)):
            # 调用汽轮机模拟
            Ds, pt, ht, D1, p1, D1j, h1, D2, p2, Dj2, h2, P2, T2, D2r, \
                D3, p3, h3, Dj3, D4, p4, h4, Dj4, D5, p5, h5, Dj5, D6, p6, \
                h6, Dj6, D7, p7, h7, Dj7, hc, Dc, P1, T1, Tj, Ty2, Ne, \
                Cond_T, Ne8 = turbine.simulation(
                Sup_wat_T, Pt_DT, Dsup_DT, Ps1, Ps2, 0, Ne_DT, P1, P2,
                Pre_real, T1, T2, Re_wat_T, Tj, Ty2, Ps3, Ps4, Ps5, Ps6, Ps7,
                Pc_sim, h
            )

            # 使用当前 kc_current 更新凝汽器模拟
            deta_Pc, Pc_sim, Dec, Tc, Hc1, Vf, NTU, Qc = coldsystem.simulation(
                hc, Dc, Pc_sim, n, Ta, h, kc_current
            )
            # Pc_sim 为凝汽器模拟得到的背压（单位KPa）

            (Ps7, Dd7, hd7, Tn7, Tl7, Tw7, Ps6, Dd6, hd6, Tn6, Tl6, Tw6,
             Ps5, Dd5, hd5, Tn5, Tl5, Tw5, Ps4, Tw4, Ps3, Dd3, hd3, Tn3,
             Tl3, Tw3, Ps2, Dd2, hd2, Tn2, Tl2, Tw2, Ps1, Dd1, hd1, Tn1,
             Tl1, Tw1, rhol4, Twc_m, Tw4_m, rhol4_m, drhol4_dt) = heatsystem.simulation(
                Cond_wat_T, Tn7, Tl7, Ps7, Tw7, D7, h, hd7, Dd6, hd6,
                Tn6, Tl6, Ps6, Tw6, D6, h, Dd5, hd5, Tn5, Tl5, Ps5, Tw5, D5, h,
                Ps4, Dd3, hd3, D4, h, Dsup_DT, Tn3, Tl3, Ps3, Tw3, D3, h,
                Dd2, hd2, Tn2, Tl2, Ps2, Tw2, D2, h, Dd1, hd1, Tn1, Tl1, Ps1, Tw1,
                D1, h, Cond_wat_D, Feed_wat_D, Tw4, Twc_m, drhol4_dt, rhol4_m, h
            )

            state_result_new = [
                Ne, Ps1, Ps2, P1, P2, T1, T2, Tj, Ty2, Ps3, Ps4, Ps5, Ps6, Ps7,
                Tn7, Tl7, Tw7, hd7, Dd6, hd6, Tn6, Tl6, Tw6, Dd5, hd5, Tn5,
                Tl5, Tw5, Dd3, hd3, Tn3, Tl3, Tw3, Dd2, hd2, Tn2, Tl2, Tw2,
                Dd1, hd1, Tn1, Tl1, Tw1, Twc_m, Tw4, drhol4_dt, rhol4_m, Pc_sim
            ]
        if not (0 in state_result_new or any(math.isnan(x) for x in state_result_new)):
            state_result = state_result_new
    result_Pc_m1 = state_result[-1] * 1000  # 计算背压

    # 判断误差是否满足条件
    Ne_m1 = state_result[0]  # 负荷
    power_consumption_fan_m1, net_power_m1 = Net_power(float(Ne_m1), m1, data1[-1, 16])
    # return result_Pc_m1, Ne_m1, power_consumption_fan_m1, net_power_m1
    return [
                    round(power_consumption_fan_m1, 2),
                    round(net_power_m1, 2),
                    round(Ne_m1, 2),
                    round(m1, 2),  # 假设 data1[11] 是风机转速
                    round(data1[-1][18], 2),  # 假设 data1[10] 是环境温度
                    round(result_Pc_m1, 2),  # 转换单位为 kPa
                    pb_list,
                    ne_list,
                    state_result
    ]


def parallel_optimize(data, simtime, sample_time=60, h=0.1, tolerance=0.0005, precision=1, logger=None):
    ita_current, kc_current, sim_steps = share_data()
    #data_start = data[-30:].copy()
    # print(f'寻优时模型汽轮机效率{ita_current}，流量系数{kc_current}')
    # data1 = data.copy()
    # print(f'寻优算法输入数据{data1}')
    # results = optimize_for_speed(data1, data1[11], simtime, h, sample_time, tolerance, ita_current, kc_current)
    # print(f'实际数据仿真到稳态后负荷{results[2]}，背压{results[5]}')
    # print(f'仿真到稳态后最终状态：{results[8]}')
    # # 显示图形
    save_path = "pics/"
    os.makedirs(save_path, exist_ok=True)
    #
    # plt.figure(figsize=(12, 8))
    # plt.title('实际背压仿真到稳态', fontsize=16)
    #
    # tick = [i for i in range(0, len(results[6]), 1)]
    # # 绘制风机功耗增量
    # plt.plot(tick, results[6], marker='o', linestyle='-', color='purple', )
    #
    # # 设置坐标轴标签
    # plt.xlabel('仿真步数', fontsize=14)
    # plt.ylabel('背压', fontsize=14)
    # plt.xticks(fontsize=12)
    # plt.yticks(fontsize=12)
    #
    # # 显示图形
    # plt.tight_layout()
    # plt.savefig(f"{save_path}/实际背压仿真到稳态.png", dpi=300)
    # plt.close()
    #
    # # TODO：最优风机转速下，负荷仿真到稳态，排汽焓的变化，等等中间变量的变化
    # plt.figure(figsize=(12, 8))
    # plt.title('实际背压仿真到稳态过程中负荷变化', fontsize=16)
    #
    # tick = [i for i in range(0, len(results[7]), 1)]
    # # 绘制风机功耗增量
    # plt.plot(tick, results[7], marker='o', linestyle='-', color='purple', )
    #
    # # 设置坐标轴标签
    # plt.xlabel('仿真步数', fontsize=14)
    # plt.ylabel('负荷', fontsize=14)
    # plt.xticks(fontsize=12)
    # plt.yticks(fontsize=12)
    #
    # # 显示图形
    # plt.tight_layout()
    # plt.savefig(f"{save_path}/实际背压仿真到稳态过程中负荷变化.png", dpi=300)
    # plt.close()
    # stable_state = results[-1]
    min_speed, max_speed = 30, 122  # 风机转速的范围
    with Pool(processes=2) as pool:  # 开启2个进程
        trace_list = []
        lower_bound, upper_bound = min_speed, max_speed  # 初始区间为[44, 122]
        best_result = None  # 用于存储与最大净功率对应的完整计算结果
        best_net_power = -np.inf  # 初始化最佳净功率为负无穷
        iterations = 0  # 初始化二分法迭代次数
        # first_bounds_saved = False  # 标记是否已保存首次区间小于等于5的上下限
        # first_lower, first_upper = None, None  # 存储首次区间小于等于5的上下限
        start_t = time.time()
        while (upper_bound - lower_bound) > precision:  # 当区间宽度大于设定精度时继续
            iterations += 1  # 增加迭代次数
            # 计算两个中点
            m1 = (lower_bound + upper_bound) / 2
            m2 = (lower_bound + m1) / 2
            logger.info(f'iteration:{iterations} | m1:{m1} m2:{m2} lower_bound:{lower_bound} upper_bound:{upper_bound}')
            data1 = data.copy()
            data2 = data.copy()
            # state1 = stable_state.copy()
            # state2 = stable_state.copy()

            results = pool.starmap(optimize_for_speed, [
                (data1, m1, simtime, h, sample_time, tolerance, ita_current, kc_current),
                (data2, m2, simtime, h, sample_time, tolerance, ita_current, kc_current)
            ])
            # results = pool.starmap(optimize_for_speed, [
            #     (data1, m1, simtime, h, sample_time, tolerance, ita_current, kc_current,state1),
            #     (data2, m2, simtime, h, sample_time, tolerance, ita_current, kc_current,state2)
            # ])
            results_wolist = [results[0][:-3],results[1][:-3]]
            trace_list = trace_list + results_wolist
            # 判断哪个中点的净功率更大
            net_power_m1 = results[0][1]
            net_power_m2 = results[1][1]
            if net_power_m1 > net_power_m2:
                # 如果 m1 的净功率更大，最大值在 [m2, upper_bound] 区间
                lower_bound = m2
                current_best_net_power = net_power_m1
                current_best_result = results[0]
            else:
                # 如果 m2 的净功率更大，最大值在 [lower_bound, m1] 区间
                upper_bound = m1
                current_best_net_power = net_power_m2
                current_best_result = results[1]

            # 更新最佳结果
            if current_best_net_power > best_net_power:
                best_net_power = current_best_net_power
                best_result = current_best_result.copy()
        end_t = time.time()
        logger.info(f'二分法寻优消耗时间{end_t - start_t}')
        # 记录与最大净功率对应的结果和迭代次数
        if best_result is not None:
            logger.info(f"已找到最大净功率，二分法迭代次数: {iterations}")
            # 将迭代次数和首次区间上下限添加到结果列表
            best_result.extend([
                iterations
            ])
            plt.figure(figsize=(12, 8))
            plt.title('背压-净功率曲线', fontsize=16)
            # plt.title('背压-净功率曲线', fontsize=16)
            # 绘制风机功耗增量
            result_array = np.array(trace_list)
            result_array = result_array[result_array[:, 5].argsort()]
            plt.plot(result_array[:, 5], result_array[:, 1], marker='o', linestyle='-',
                     color='purple', )

            # 设置坐标轴标签
            plt.xlabel('背压', fontsize=14)
            plt.ylabel('净功率', fontsize=14)
            plt.xticks(fontsize=12)
            plt.yticks(fontsize=12)

            # 显示图形
            plt.tight_layout()
            plt.savefig(f"{save_path}/背压-净功率曲线.png", dpi=300)
            plt.close()

            plt.figure(figsize=(12, 8))
            plt.title('风机转速-背压曲线', fontsize=16)
            # plt.title('背压-净功率曲线', fontsize=16)
            # 绘制风机功耗增量
            result_array = np.array(trace_list)
            result_array = result_array[result_array[:, 3].argsort()]
            plt.plot(result_array[:, 3], result_array[:, 5], marker='o', linestyle='-',
                     color='purple', )

            # 设置坐标轴标签
            plt.xlabel('风机转速', fontsize=14)
            plt.ylabel('背压', fontsize=14)
            plt.xticks(fontsize=12)
            plt.yticks(fontsize=12)

            # 显示图形
            plt.tight_layout()
            plt.savefig(f"{save_path}/风机转速-背压曲线.png", dpi=300)
            plt.close()

            plt.figure(figsize=(12, 8))
            plt.title('背压-负荷曲线', fontsize=16)
            # plt.title('背压-净功率曲线', fontsize=16)
            # 绘制风机功耗增量
            result_array = np.array(trace_list)
            result_array = result_array[result_array[:, 5].argsort()]
            plt.plot(result_array[:, 5], result_array[:, 2], marker='o', linestyle='-',
                     color='purple', )

            # 设置坐标轴标签
            plt.xlabel('背压', fontsize=14)
            plt.ylabel('负荷', fontsize=14)
            plt.xticks(fontsize=12)
            plt.yticks(fontsize=12)

            # 显示图形
            plt.tight_layout()
            plt.savefig(f"{save_path}/背压-负荷曲线.png", dpi=300)
            plt.close()

            plt.figure(figsize=(12, 8))
            plt.title('最优背压仿真到稳态', fontsize=16)

            tick = [i for i in range(0, len(best_result[6]), 1)]
            # 绘制风机功耗增量
            plt.plot(tick, best_result[6], marker='o', linestyle='-', color='purple', )

            # 设置坐标轴标签
            plt.xlabel('仿真步数', fontsize=14)
            plt.ylabel('背压', fontsize=14)
            plt.xticks(fontsize=12)
            plt.yticks(fontsize=12)

            # 显示图形
            plt.tight_layout()
            plt.savefig(f"{save_path}/最优背压仿真到稳态.png", dpi=300)
            plt.close()

            # TODO：最优风机转速下，负荷仿真到稳态，排汽焓的变化，等等中间变量的变化
            plt.figure(figsize=(12, 8))
            plt.title('最优背压仿真到稳态过程中负荷变化', fontsize=16)

            tick = [i for i in range(0, len(best_result[7]), 1)]
            # 绘制风机功耗增量
            plt.plot(tick, best_result[7], marker='o', linestyle='-', color='purple', )

            # 设置坐标轴标签
            plt.xlabel('仿真步数', fontsize=14)
            plt.ylabel('负荷', fontsize=14)
            plt.xticks(fontsize=12)
            plt.yticks(fontsize=12)

            # 显示图形
            plt.tight_layout()
            plt.savefig(f"{save_path}/最优背压仿真到稳态过程中负荷变化.png", dpi=300)
            plt.close()
            #
            with open('opt_data.pkl', 'wb') as f:
                pickle.dump({'result_array': result_array, 'best_result': best_result}, f)
            #np.save('result_array_last.npy', result_array)
            #np.save('best_result_last.npy', best_result)
            # np.savez('opt_arrays.npz', result_array=result_array, best_result=best_result)

    print(f'优化后负荷：{np.array(best_result[:-4],np.float32)[[2]]}')
    print(f'优化后风机功耗：{np.array(best_result[:-4], np.float32)[[0]]}')

    return np.array(best_result[:-4],np.float32)[[3,5,1]], trace_list